{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI;\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    api_key='sk-fumkL2Nvub7ibL0J3d27237a09Df4e3194EcE6E317149a60',\n",
    "    base_url='https://api.gpts.vin/v1'\n",
    "    )\n",
    "\n",
    "\n",
    "print(llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.llms import Ollama\n",
    "\n",
    "# llm = Ollama(model='llama3')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Schema\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class Person(BaseModel):\n",
    "    \"\"\"Information about a person\"\"\"\n",
    "\n",
    "    name: Optional[str] = Field(default=None, description=\"The name of the person\")\n",
    "    hair_color: Optional[str] = Field(default=None, description=\"The color of the person hair if known\")\n",
    "    height_in_meters: Optional[str] = Field(default=None, description=\"Height measured in meters\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Extractor\n",
    "\n",
    "Let's create an information extractor using the schema we defined above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "# Define a custom prompt to provide instructions and any additonal context\n",
    "# 1) You can add examples into the prompt template to improve extraction quality\n",
    "# 2) Introduce additonal parameters to take context into account\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are an expert extraction algorithm.\"),\n",
    "        (\"human\", \"{text}\")\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_mistralai import ChatMistralAI\n",
    "\n",
    "# llm = ChatMistralAI(model=\"mistral-large-latest\", temperature=0)\n",
    "\n",
    "runnable = prompt | llm.with_structured_output(schema=Person)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"wangxijun is 6 feet tall and has blond hair.\"\n",
    "\n",
    "runnable.invoke({\"text\": text})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple Entities\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "class Data(BaseModel):\n",
    "    \n",
    "    # Create a model so that we can extract multiple entities\n",
    "    people: List[Person]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runnable = prompt | llm.with_structured_output(schema=Data)\n",
    "\n",
    "text = \"My name is Jeff, my hair is black and i am 6 feet tall. Anna has the same color hair as me.\"\n",
    "\n",
    "runnable.invoke({\"text\": \"{text}\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
